"""
读取ACL IMDB数据集信息
"""
import os
import sys
import tarfile

import pandas as pd
from tqdm import tqdm

from config.sys_config import ACL_IMDB_DATASET_PATH
from dataset._base_uitls import download_dataset


ACL_IMDB_DATASET_URL = 'http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz'


def read_acl_imdb(download=False, chunk_size=1024 * 10) -> tarfile.TarFile:
    if not os.path.exists(ACL_IMDB_DATASET_PATH):
        if download:
            download_dataset(ACL_IMDB_DATASET_PATH, ACL_IMDB_DATASET_URL, chunk_size)
        else:
            print("Error: Dataset is not exist.Please use arg `download=True` to start download.")
            return sys.exit(1)
    return tarfile.open(ACL_IMDB_DATASET_PATH)


class AclImdbDataset:
    _ACL_IMDB_FILE: tarfile.TarFile
    _FILENAME_DATA_FRAME: pd.DataFrame
    _TRAIN_DATA_FRAME: pd.DataFrame
    _TRAIN_UNSUP_DATA_FRAME: pd.DataFrame
    _TEST_DATA_FRAME: pd.DataFrame
    _TRAIN_NEG_PATH: str = 'aclImdb/train/neg'
    _TRAIN_POS_PATH: str = 'aclImdb/train/pos'
    _TRAIN_UNSUP_PATH: str = 'aclImdb/train/unsup'
    _TEST_NEG_PATH: str = 'aclImdb/test/neg'
    _TEST_POS_PATH: str = 'aclImdb/test/pos'

    def __init__(self, tarFile: tarfile.TarFile = None, download=True, chunk_size=1024 * 10):
        if tarFile is None:
            self._ACL_IMDB_FILE = read_acl_imdb(download, chunk_size)
        else:
            self._ACL_IMDB_FILE = tarFile
        self._read_file_list()
        self._create_data_frame()

    def _create_data_frame(self):
        self._TRAIN_DATA_FRAME = self._FILENAME_DATA_FRAME.loc[
            self._FILENAME_DATA_FRAME['path'].isin([self._TRAIN_NEG_PATH, self._TRAIN_POS_PATH])].copy()
        self._TEST_DATA_FRAME = self._FILENAME_DATA_FRAME.loc[
            self._FILENAME_DATA_FRAME['path'].isin([self._TEST_NEG_PATH, self._TEST_POS_PATH])].copy()
        self._TRAIN_UNSUP_DATA_FRAME = self._FILENAME_DATA_FRAME.loc[
                    self._FILENAME_DATA_FRAME['path'] == self._TRAIN_UNSUP_PATH].copy()

        self._TRAIN_DATA_FRAME.sort_values(by=['label', 'file_id'], ascending=[True, True], inplace=True)
        self._TEST_DATA_FRAME.sort_values(by=['label', 'file_id'], ascending=[True, True], inplace=True)
        self._TRAIN_UNSUP_DATA_FRAME.sort_values(by=['label', 'file_id'], ascending=[True, True], inplace=True)

    def _read_file_list(self):
        _filename_list = []
        for member in tqdm(self._ACL_IMDB_FILE.getmembers(), desc="正在读取文件信息"):
            f = self._ACL_IMDB_FILE.extractfile(member)
            if f is not None and not member.name.endswith('.feat'):
                if member.name.startswith(self._TRAIN_NEG_PATH) \
                        or member.name.startswith(self._TRAIN_POS_PATH) \
                        or member.name.startswith(self._TRAIN_UNSUP_PATH) \
                        or member.name.startswith(self._TEST_NEG_PATH) \
                        or member.name.startswith(self._TEST_POS_PATH):
                    file_context = f.read().decode('utf-8')
                    path, filename = member.name.rsplit('/', 1)
                    file_id, ranking = filename.replace('.txt', '').split('_')
                    row = (path.rsplit('/')[-1], int(file_id), int(ranking), filename, path, file_context)
                    _filename_list.append(row)
        self._FILENAME_DATA_FRAME = pd.DataFrame(_filename_list,
                                                 columns=['label', 'file_id', 'ranking', 'filename', 'path', 'content'])
        self._ACL_IMDB_FILE.close()

    @property
    def train_data(self) -> pd.DataFrame:
        return self._TRAIN_DATA_FRAME.copy()

    @property
    def train_data_size(self) -> int:
        return len(self._TRAIN_DATA_FRAME)

    @property
    def test_data(self) -> pd.DataFrame:
        return self._TEST_DATA_FRAME.copy()

    @property
    def test_data_size(self) -> int:
        return len(self._TEST_DATA_FRAME)

    @property
    def file_data(self) -> pd.DataFrame:
        return self._FILENAME_DATA_FRAME.copy()

    @property
    def file_data_size(self) -> int:
        return len(self._FILENAME_DATA_FRAME)

    @property
    def unsup_data(self) -> pd.DataFrame:
        return self._TRAIN_UNSUP_DATA_FRAME.copy()

    @property
    def unsup_data_size(self) -> int:
        return len(self._TRAIN_UNSUP_DATA_FRAME)


if __name__ == '__main__':
    acl_imdb_file = read_acl_imdb(download=False)
    dataset = AclImdbDataset(tarFile=acl_imdb_file)
    print(dataset.train_data[0:10][['label', 'file_id', 'ranking', 'path']])
    print(dataset.test_data[0:10][['label', 'file_id', 'ranking', 'path']])
    print(dataset.unsup_data[0:10][['label', 'file_id', 'ranking', 'path']])
